from collections import defaultdict, Counter
import networkx as nx
import community
__author__ = 'panagiotis'


class SpreadActivation(nx.DiGraph):  # maybe not digraph but graph -> not predecessors or successors but neighbors

    def __init__(self, G=None, spreading_factor=0.25, accuracy_threshold=0.01, activation_threshold=0.0):
        if not G:
            nx.DiGraph.__init__(self)
        elif isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
            nx.DiGraph.__init__(self, G)
        elif isinstance(G, (nx.Graph, nx.MultiGraph)):
            nx.DiGraph.__init__(self, G.to_directed())
        else:
            raise TypeError("Parameter G should be a networkx Graph or subclass.")
        self.spreading_factor = spreading_factor
        self.accuracy_threshold = accuracy_threshold
        self.activation_threshold = activation_threshold
        self.deviation = defaultdict(float)
        self.active = set()
        self.iterations = 0

    def setup_energy(self, active, initial_energy, energy="energy"):
        self.active |= set(active)
        for (x, y) in self.edges():
            self.edge[x][y][energy] = 0.0
        if hasattr(initial_energy, "__iter__"):
            for (u, e) in zip(active, initial_energy):
                self.node[u][energy] = e
            for u in set(self.nodes()) - set(active):
                self.node[u][energy] = 0.0
        else:
            for x in self.nodes():
                if x in active:
                    self.node[x][energy] = initial_energy
                    self.node[x]["prv"] = initial_energy
                else:
                    self.node[x][energy] = 0.0
                    self.node[x]["prv"] = 0.0

    def normalize_weights(self, strength="weight"):
        import math
        for (x, y, d) in self.edges(data=True):
            # s = (2 / (1+math.exp(-self.edge[x][y][strength])) - 1)
            s = math.log(self.edge[x][y][strength]+1)+1
            if s > 0:
                self.edge[x][y][strength] = s
            else:
                self.remove_edge(x, y)

    def predict(self, spreading_factor=None, accuracy_threshold=None, activation_threshold=None,
                strength="weight", energy="energy", max_iter=15):
        if spreading_factor: self.spreading_factor = spreading_factor
        if accuracy_threshold: self.accuracy_threshold = accuracy_threshold
        if activation_threshold: self.activation_threshold = activation_threshold
        self.iterations = 0
        while True:
            self.iterations += 1
            condition1 = set([x for (x, d) in self.nodes(data=True) if d[energy] > 0]) - self.active == set()
            self.active = set([x for (x, d) in self.nodes(data=True) if d[energy] > 0])
            # print "iteration", self.iterations
            # for every active node X
            edge_total_energy = 0
            for x in self.active:
                energy_to_spread = self.spreading_factor * self.node[x][energy]
                edge_total_energy += energy_to_spread
                # calculate the energy node X retains.
                self.node[x][energy] -= energy_to_spread
                # for every successor Y that node X has, (edge from X to Y exists)
                spread_energy = 0.0
                denominator = sum([self.edge[x][s][strength] for s in self.successors(x)])
                for y in self.successors(x):
                    # calculate the energy to transfer from X to Y, depending to the relative tie strength of X to Y,
                    # compared to the sum of ties strength of all outgoing edges of X.
                    tie_strength = self.edge[x][y][strength]
                    try:
                        self.edge[x][y][energy] = ((energy_to_spread * tie_strength) / denominator)
                                                   # self.out_degree(x, weight=strength)))
                        spread_energy += self.edge[x][y][energy]
                    except ZeroDivisionError:   # this means that all ties have 0.0 strength
                        self.edge[x][y][energy] = 0.0
                        print "ZeroDivisionError for edge", x, y, ", so edge was set with 0.0 energy"
                self.node[x][energy] += energy_to_spread - spread_energy
            # for every node
            for x in self.nodes():
                # store energy node retained.
                energy_retained = self.node[x][energy]
                # calculate energy node has received.
                energy_received = 0.0
                for y in self.predecessors(x):
                    try:
                        energy_received += self.edge[y][x][energy]
                    except ValueError:
                        self.node[x][energy] += 0.0
                        print "ValueError for edge", y, x, "has no energy key"
                    except KeyError:
                        print "KeyError for edge", y, x
                    finally:
                        self.edge[y][x][energy] = 0.0
                # energy_received += self.in_degree(x, weight=energy)
                # energy node has is sum of energy retained and energy received
                self.node[x][energy] = energy_retained + energy_received
                self.deviation[x] = abs(self.node[x]['prv'] - self.node[x][energy])
                self.node[x]['prv'] = self.node[x][energy]
            condition2 = all([self.deviation[x] < self.accuracy_threshold for x in self.active])
            # print Counter(self.deviation).most_common(5)
            if condition1 and condition2 or self.iterations >= max_iter:
                # print self.iterations, "iterations to converge"
                return self.most_common()

    def most_common(self, n=None, energy="energy", percent=100):
        from collections import Counter
        results = Counter()
        for x in self.nodes():
            results[x] = self.node[x][energy]
        if not n:
            n = int((percent / 100.0) * self.number_of_nodes())
        return results.most_common(n)

    def above_threshold(self, threshold=1.0, energy="energy"):
        from collections import Counter
        results = Counter()
        for x in self.nodes():
            results[x] = self.node[x][energy]
        return [u for (u, p) in results.most_common() if p >= threshold]


class CommunityActivation(nx.DiGraph):

    def __init__(self, G, strength="strength"):
        T = nx.DiGraph()
        if isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
            predicted = community.best_partition(G.to_undirected())
        elif isinstance(G, (nx.Graph, nx.MultiGraph)):
            predicted = community.best_partition(G)
        else:
            raise TypeError("Parameter G should be a networkx Graph or subclass.")
        for com in [[n for n in G.nodes() if predicted[n] == v] for v in set(predicted.values())]:
            T.add_node(SpreadActivation(G.subgraph(com)))
        for oC in T.nodes():
            for tC in [t for t in T.nodes() if t != oC]:
                T.add_edge(oC, tC, strength=sum([sum([G[n1][n2][strength] for n2 in G.successors(n1) if n2 in tC.nodes()]) for n1 in oC.nodes()]))
        nx.DiGraph.__init__(self, T)

    def setup_energy(self, active, initial_energy, energy="energy"):
        if hasattr(initial_energy, "__iter__"):
            for G in self.nodes():
                g_active, g_initial = zip(*[(u, e) for (u, e) in zip(active, initial_energy) if u in G.nodes()])
                G.setup_energy(g_active, g_initial, energy)
        else:
            for G in self.nodes():
                G.setup_energy((set(G.nodes()) & set(active)), initial_energy, energy)

    def normalize_weights(self, strength="weight"):
        for G in self.nodes():
            G.normalize_weights(strength)

    def predict(self, spreading_factor=None, accuracy_threshold=None, activation_threshold=None,
                strength="weight", energy="energy", max_iter=15):
        results = []
        for G in self.nodes():
            results += G.predict(spreading_factor, accuracy_threshold, activation_threshold,
                                 strength, energy, max_iter)
        return results

    def most_common(self, n=None, energy="energy", percent=100):
        aggregate = []
        for G in self.nodes():
            aggregate += G.most_common()
        from collections import Counter
        results = Counter()
        for (p, e) in aggregate:
            results[p] = e
        if not n:
            n = int((percent / 100.0) * len(results))
        return results.most_common(n)


class UndirectedSpreadActivation(nx.Graph):  # maybe not digraph but graph -> not predecessors or successors but neighbors

    def __init__(self, G=None, spreading_factor=0.25, accuracy_threshold=0.01, activation_threshold=0.0):
        if not G:
            nx.Graph.__init__(self)
        elif isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
            nx.Graph.__init__(self, G.to_undirected())
        elif isinstance(G, (nx.Graph, nx.MultiGraph)):
            nx.Graph.__init__(self, G)
        else:
            raise TypeError("Parameter G should be a networkx Graph or subclass.")
        self.spreading_factor = spreading_factor
        self.accuracy_threshold = accuracy_threshold
        self.activation_threshold = activation_threshold
        self.deviation = defaultdict(float)
        self.active = set()
        self.iterations = 0

    def setup_energy(self, active, initial_energy, energy="energy"):
        self.active |= set(active)
        for (x, y) in self.edges():
            self.edge[x][y][energy] = 0.0
        if hasattr(initial_energy, "__iter__"):
            for (u, e) in zip(active, initial_energy):
                self.node[u][energy] = e
            for u in set(self.nodes()) - set(active):
                self.node[u][energy] = 0.0
        else:
            for x in self.nodes():
                if x in active:
                    self.node[x][energy] = initial_energy
                    self.node[x]["prv"] = initial_energy
                else:
                    self.node[x][energy] = 0.0
                    self.node[x]["prv"] = 0.0

    def normalize_weights(self, strength="weight"):
        import math
        for (x, y, d) in self.edges(data=True):
            # s = (2 / (1+math.exp(-self.edge[x][y][strength])) - 1)
            s = math.log(self.edge[x][y][strength]+1)+1
            if s > 0:
                self.edge[x][y][strength] = s
            else:
                self.remove_edge(x, y)

    def predict(self, spreading_factor=None, accuracy_threshold=None, activation_threshold=None,
                strength="weight", energy="energy", max_iter=15):
        if spreading_factor: self.spreading_factor = spreading_factor
        if accuracy_threshold: self.accuracy_threshold = accuracy_threshold
        if activation_threshold: self.activation_threshold = activation_threshold
        self.iterations = 0
        while True:
            self.iterations += 1
            condition1 = set([x for (x, d) in self.nodes(data=True) if d[energy] > 0]) - self.active == set()
            self.active = set([x for (x, d) in self.nodes(data=True) if d[energy] > 0])
            # print "iteration", self.iterations
            # for every active node X
            edge_total_energy = 0
            for x in self.active:
                energy_to_spread = self.spreading_factor * self.node[x][energy]
                edge_total_energy += energy_to_spread
                # calculate the energy node X retains.
                self.node[x][energy] -= energy_to_spread
                # for every successor Y that node X has, (edge from X to Y exists)
                spread_energy = 0.0
                denominator = sum([self.edge[x][s][strength] for s in self.neighbors(x)])
                for y in self.neighbors(x):
                    # calculate the energy to transfer from X to Y, depending to the relative tie strength of X to Y,
                    # compared to the sum of ties strength of all outgoing edges of X.
                    tie_strength = self.edge[x][y][strength]
                    try:
                        self.edge[x][y][energy] = ((energy_to_spread * tie_strength) / denominator)
                                                   # self.out_degree(x, weight=strength)))
                        spread_energy += self.edge[x][y][energy]
                    except ZeroDivisionError:   # this means that all ties have 0.0 strength
                        self.edge[x][y][energy] = 0.0
                        print "ZeroDivisionError for edge", x, y, ", so edge was set with 0.0 energy"
                self.node[x][energy] += energy_to_spread - spread_energy
            # for every node
            for x in self.nodes():
                # store energy node retained.
                energy_retained = self.node[x][energy]
                # calculate energy node has received.
                energy_received = 0.0
                for y in self.neighbors(x):
                    try:
                        energy_received += self.edge[y][x][energy]
                    except ValueError:
                        self.node[x][energy] += 0.0
                        print "ValueError for edge", y, x, "has no energy key"
                    except KeyError:
                        print "KeyError for edge", y, x
                    finally:
                        self.edge[y][x][energy] = 0.0
                # energy_received += self.in_degree(x, weight=energy)
                # energy node has is sum of energy retained and energy received
                self.node[x][energy] = energy_retained + energy_received
                self.deviation[x] = abs(self.node[x]['prv'] - self.node[x][energy])
                self.node[x]['prv'] = self.node[x][energy]
            condition2 = all([self.deviation[x] < self.accuracy_threshold for x in self.active])
            # print Counter(self.deviation).most_common(5)
            if condition1 and condition2 or self.iterations >= max_iter:
                # print self.iterations, "iterations to converge"
                return self.most_common()

    def most_common(self, n=None, energy="energy", percent=100):
        from collections import Counter
        results = Counter()
        for x in self.nodes():
            results[x] = self.node[x][energy]
        if not n:
            n = int((percent / 100.0) * self.number_of_nodes())
        return results.most_common(n)

    def above_threshold(self, threshold=1.0, energy="energy"):
        from collections import Counter
        results = Counter()
        for x in self.nodes():
            results[x] = self.node[x][energy]
        return [u for (u, p) in results.most_common() if p >= threshold]

